"""
Hybrid PETSc-Kokkos Extension Implementation

Main extension class that integrates the Hybrid PETSc-Kokkos framework with GSM.
"""

from typing import Dict, List, Any
from pathlib import Path
import os
import logging

from gcr_solver_manager.extensions.base_extension import BaseExtension
from gcr_solver_manager.extensions.base_builder import BaseBuilder
from gcr_solver_manager.extensions.base_runner import BaseRunner
from gcr_solver_manager.extensions.base_validator import BaseValidator

from .builder import HybridBuilder
from .runner import HybridRunner
from .validator import HybridValidator

logger = logging.getLogger(__name__)


class HybridPetscKokkosExtension(BaseExtension):
    """
    Extension for Hybrid PETSc-Kokkos Integration framework.
    
    Provides support for building and running hybrid solvers that combine
    PETSc linear algebra with Kokkos performance portability.
    """
    
    def __init__(self):
        super().__init__("hybrid-petsc-kokkos", "1.0.0")
        self._builder = None
        self._runner = None
        self._validator = None
    
    @property
    def description(self) -> str:
        """Get extension description."""
        return ("Hybrid PETSc-Kokkos Integration framework providing high-performance "
                "solvers with intelligent framework dispatch and GPU acceleration")
    
    @property
    def supported_frameworks(self) -> List[str]:
        """Get supported framework names."""
        return [
            "hybrid-petsc-kokkos",
            "hybrid-gcr",
            "hybrid-ca-gcr", 
            "hybrid-gmres",
            "hybrid-bca-gmres"
        ]
    
    @property
    def config_schema(self) -> Dict[str, Any]:
        """Get JSON schema for configuration validation."""
        return {
            "$schema": "http://json-schema.org/draft-07/schema#",
            "type": "object",
            "title": "Hybrid PETSc-Kokkos Configuration Schema",
            "properties": {
                "solver": {
                    "type": "object",
                    "properties": {
                        "type": {
                            "type": "string",
                            "enum": ["gcr", "ca-gcr", "gmres", "bca-gmres"],
                            "default": "gcr"
                        },
                        "tolerance": {
                            "type": "number",
                            "minimum": 1e-12,
                            "maximum": 1e-1,
                            "default": 1e-6
                        },
                        "max_iterations": {
                            "type": "integer",
                            "minimum": 1,
                            "maximum": 10000,
                            "default": 1000
                        },
                        "restart": {
                            "type": "integer",
                            "minimum": 1,
                            "maximum": 1000,
                            "default": 50
                        }
                    },
                    "required": ["type"]
                },
                "matrix": {
                    "type": "object",
                    "properties": {
                        "file": {"type": "string"},
                        "format": {
                            "type": "string",
                            "enum": ["matrix_market", "petsc_binary", "hdf5"],
                            "default": "matrix_market"
                        },
                        "size": {"type": "integer", "minimum": 1}
                    },
                    "required": ["file"]
                },
                "execution": {
                    "type": "object",
                    "properties": {
                        "mpi_processes": {
                            "type": "integer",
                            "minimum": 1,
                            "default": 1
                        },
                        "openmp_threads": {
                            "type": "integer", 
                            "minimum": 1,
                            "default": 1
                        },
                        "gpu_devices": {
                            "type": "integer",
                            "minimum": 0,
                            "default": 0
                        },
                        "memory_limit": {"type": "string", "default": "8GB"}
                    }
                },
                "petsc": {
                    "type": "object",
                    "properties": {
                        "ksp_type": {"type": "string"},
                        "pc_type": {"type": "string", "default": "none"},
                        "options": {
                            "type": "array",
                            "items": {"type": "string"}
                        }
                    }
                },
                "kokkos": {
                    "type": "object",
                    "properties": {
                        "backend": {
                            "type": "string",
                            "enum": ["serial", "openmp", "cuda", "hip", "sycl"],
                            "default": "serial"
                        },
                        "memory_space": {"type": "string"},
                        "execution_space": {"type": "string"}
                    }
                }
            },
            "required": ["solver", "matrix"]
        }
    
    def can_handle_config(self, config_path: str) -> bool:
        """
        Check if this extension can handle the given configuration.
        
        Args:
            config_path: Path to configuration file
            
        Returns:
            True if extension can handle this configuration
        """
        config_path_obj = Path(config_path)
        
        # Check path patterns
        path_str = str(config_path_obj)
        
        # High confidence matches
        if "prototypes/hybrid" in path_str:
            logger.debug(f"Hybrid path pattern matched: {config_path}")
            return True
        
        if "hybrid_prototype" in path_str:
            logger.debug(f"Hybrid prototype pattern matched: {config_path}")
            return True
        
        # Medium confidence matches
        if config_path_obj.name.startswith("hybrid_"):
            return True
        
        # Check if config file exists and contains hybrid markers
        if config_path_obj.exists():
            try:
                content = config_path_obj.read_text().lower()
                hybrid_markers = [
                    "hybrid",
                    "petsc",
                    "kokkos", 
                    "gcr-nccl-hybrid",
                    "ca-gcr",
                    "bca-gmres"
                ]
                
                marker_count = sum(1 for marker in hybrid_markers if marker in content)
                if marker_count >= 2:  # Need at least 2 markers for confidence
                    logger.debug(f"Content analysis matched: {config_path} (markers: {marker_count})")
                    return True
                    
            except Exception as e:
                logger.debug(f"Error reading config file {config_path}: {e}")
        
        return False
    
    def get_builder(self) -> BaseBuilder:
        """Get the builder instance for this extension."""
        if self._builder is None:
            self._builder = HybridBuilder("hybrid-petsc-kokkos")
        return self._builder
    
    def get_runner(self) -> BaseRunner:
        """Get the runner instance for this extension."""
        if self._runner is None:
            self._runner = HybridRunner("hybrid-petsc-kokkos")
        return self._runner
    
    def get_validator(self) -> BaseValidator:
        """Get the validator instance for this extension.""" 
        if self._validator is None:
            self._validator = HybridValidator("hybrid-petsc-kokkos")
        return self._validator
    
    def initialize(self) -> None:
        """Initialize the extension."""
        logger.info("Initializing Hybrid PETSc-Kokkos extension")
        
        # Check environment
        validation_errors = self.validate_environment()
        if validation_errors:
            logger.warning(f"Environment validation warnings: {validation_errors}")
        
        # Initialize components
        try:
            self.get_builder()
            self.get_runner()
            self.get_validator()
            logger.info("Hybrid PETSc-Kokkos extension initialized successfully")
        except Exception as e:
            logger.error(f"Failed to initialize extension components: {e}")
            raise
    
    def cleanup(self) -> None:
        """Clean up extension resources."""
        logger.info("Cleaning up Hybrid PETSc-Kokkos extension")
        
        # Clean up component instances
        self._builder = None
        self._runner = None
        self._validator = None
    
    def validate_environment(self) -> List[str]:
        """
        Validate that the environment meets extension requirements.
        
        Returns:
            List of error messages (empty if valid)
        """
        errors = []
        
        # Check for PETSc
        petsc_dir = os.getenv('PETSC_DIR')
        if not petsc_dir:
            errors.append("PETSC_DIR environment variable not set")
        elif not Path(petsc_dir).exists():
            errors.append(f"PETSC_DIR path does not exist: {petsc_dir}")
        
        # Check for Kokkos
        kokkos_root = os.getenv('Kokkos_ROOT')
        if not kokkos_root:
            # Also check common installation paths
            common_paths = [
                Path("/usr/local/include/Kokkos_Core.hpp"),
                Path("/opt/kokkos/include/Kokkos_Core.hpp"),
            ]
            if not any(p.exists() for p in common_paths):
                errors.append("Kokkos not found (set Kokkos_ROOT or install to standard location)")
        elif not Path(kokkos_root).exists():
            errors.append(f"Kokkos_ROOT path does not exist: {kokkos_root}")
        
        # Check for MPI
        mpi_found = any(
            Path(f"/usr/bin/{cmd}").exists() or 
            Path(f"/usr/local/bin/{cmd}").exists() 
            for cmd in ["mpirun", "mpiexec"]
        )
        if not mpi_found:
            errors.append("MPI not found in standard locations")
        
        # Check for CMake
        cmake_found = any(
            Path(f"/usr/bin/cmake").exists(),
            Path(f"/usr/local/bin/cmake").exists()
        )
        if not cmake_found:
            errors.append("CMake not found in standard locations")
        
        return errors
    
    def get_commands(self) -> Dict[str, Any]:
        """Get custom CLI commands provided by this extension."""
        return {
            "hybrid": {
                "help": "Hybrid PETSc-Kokkos framework commands",
                "subcommands": {
                    "build": {
                        "help": "Build hybrid framework",
                        "args": [
                            {"name": "--clean", "action": "store_true", "help": "Clean build"},
                            {"name": "--verbose", "action": "store_true", "help": "Verbose output"},
                        ]
                    },
                    "test": {
                        "help": "Run hybrid framework tests",
                        "args": [
                            {"name": "--quick", "action": "store_true", "help": "Quick test suite"},
                            {"name": "--comprehensive", "action": "store_true", "help": "Full test suite"},
                        ]
                    },
                    "benchmark": {
                        "help": "Run performance benchmarks",
                        "args": [
                            {"name": "--compare", "action": "store_true", "help": "Compare with baseline"},
                        ]
                    }
                }
            }
        }
    
    def get_dependencies(self) -> List[str]:
        """Get list of other extensions this extension depends on."""
        return []  # No dependencies on other extensions